'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 6_study_hyde_retriever.py
* @Time: 2025/10/30
* @All Rights Reserve By Brtc
'''
import dotenv
import weaviate
from langchain.chains import hyde
from langchain_community.chat_models.writer import ChatWriter
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.language_models import BaseLanguageModel
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.retrievers import BaseRetriever
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_weaviate import WeaviateVectorStore

dotenv.load_dotenv()

class HyDeRetriever(BaseRetriever):
    """HyDe混合检索策略"""
    retriever: BaseRetriever
    llm:BaseLanguageModel

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> list[Document]:
        """传递检索的query 实现 Hyde混合检索策略"""
        #1、构建生成假设性文档的prompt
        prompt = ChatPromptTemplate.from_template(""" 
            请写一篇科学论文来回答这个问题.\n
            问题是:{question}\n
            文章:
        """)
        #2、构建链
        chain = (
            {"question":RunnablePassthrough()}
            |prompt
            |self.llm
            |StrOutputParser()
            |self.retriever
        )
        return chain.invoke(query)

client = weaviate.connect_to_local("192.168.106.129", 8080)
db = WeaviateVectorStore(
    client,
    index_name="TestDemo",
    text_key="text",
    embedding=OpenAIEmbeddings(model="text-embedding-3-small")
)
retriever = db.as_retriever(search_type="mmr")

# 创建Hyde 检索器
hyde_retriever = HyDeRetriever(
    retriever=retriever,
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
)

# 检索文档
docs = hyde_retriever.invoke("关于LLMOPS的应用配置有哪些？")
for doc in docs:
    print("=========================================")
    print(doc.page_content[:50])
client.close()